#include <sophus/se3.hpp>
#include <string>
#include <iostream>
#include <fstream>
#include <unistd.h>

// need pangolin for plotting trajectory
#include <pangolin/pangolin.h>

using namespace std;
using namespace Eigen;

// path to trajectory file
string str_estimated_trajectory_file = "../estimated.txt";
string str_groundtruth_trajectory_file = "../groundtruth.txt";

// function for plotting trajectory, don't edit this code
// start point is red and end point is blue
void DrawTrajectory(vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_est,
	vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_gt);

 
/*******************************计算轨迹误差*********************************************/
double ErrorTrajectory(vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_g,
        vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_e );

int main(int argc, char **argv) {

    vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> estimated_poses;
    vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> groundtruth_poses;

    /// implement pose reading code
    // start your code here (5~10 lines)
    ifstream fin_est(str_estimated_trajectory_file);
    if(!fin_est)
    {
        cout<<"Cannot find estimated trajectory file at "<< str_estimated_trajectory_file << endl;
        cout<<"Exit this program!!!"<<endl;
        return -1;
    }
    else
    {
        cout<<"Have find estimated trajectory file and reading it ..."<<endl; 
        while(!fin_est.eof())
        {
            double time, tx, ty, tz, qx, qy, qz, qw;
            fin_est>> time >> tx >>ty>> tz >> qx >> qy >> qz >> qw;
            Quaterniond Quat(qw, qx, qy, qz);
            Quat.normalized();
            Vector3d Trans(tx,ty,tz);
            Sophus::SE3<double> SE3_Twr(Quat, Trans);
            estimated_poses.push_back(SE3_Twr);
        }
    }

    ifstream fin_gt(str_groundtruth_trajectory_file);
    if(!fin_gt)
    {
        cout<<"Cannot find groundtruth trajectory file at "<< str_groundtruth_trajectory_file << endl;
        cout<<"Exit this program!!!"<<endl;
        return -1;
    }
    else
    {
        cout<<"Have find groundtruth trajectory file and reading it ..."<<endl; 
        while(!fin_gt.eof())
        {
            double time, tx, ty, tz, qx, qy, qz, qw;
            fin_gt>> time >> tx >>ty>> tz >> qx >> qy >> qz >> qw;
            Quaterniond Quat(qw, qx, qy, qz);
            Quat.normalized();
            Vector3d Trans(tx,ty,tz);
            Sophus::SE3<double> SE3_Twr(Quat, Trans);
            groundtruth_poses.push_back(SE3_Twr);
        }
    }

    double rmse_value = ErrorTrajectory(groundtruth_poses, estimated_poses);
    cout<< "rmse value: "<< rmse_value << endl;
    // end your code here

    // draw trajectory in pangolin
    DrawTrajectory(estimated_poses, groundtruth_poses);
    return 0;
}

 
/*******************************计算轨迹误差*********************************************/
double ErrorTrajectory(vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_g,
        vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_e )
{
    double RMSE = 0;
    Matrix<double ,6,1> se3;
    vector<double> error;
    for(int i=0;i<poses_g.size();i++){
        se3=(poses_g[i].inverse()*poses_e[i]).log();  //这里的se3为向量形式，求log之后是向量形式
        //cout<<se3.transpose()<<endl;
        error.push_back( se3.squaredNorm() );  //二范数
       // cout<<error[i]<<endl;
    }
 
    for(int i=0; i<poses_g.size();i++){
        RMSE += error[i];
    }
    RMSE /= double(error.size());
    RMSE = sqrt(RMSE);
    return RMSE;
}

/*******************************************************************************************/
void DrawTrajectory(vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_est,
	vector<Sophus::SE3<double>, Eigen::aligned_allocator<Sophus::SE3<double>>> poses_gt) {
    if (poses_est.empty() || poses_gt.empty()) {
        cerr << "Trajectory is empty!" << endl;
        return;
    }

    // create pangolin window and plot the trajectory
    pangolin::CreateWindowAndBind("Trajectory Viewer", 1024, 768);
    glEnable(GL_DEPTH_TEST);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    pangolin::OpenGlRenderState s_cam(
            pangolin::ProjectionMatrix(1024, 768, 500, 500, 512, 389, 0.1, 1000),
            pangolin::ModelViewLookAt(0, -0.1, -1.8, 0, 0, 0, 0.0, -1.0, 0.0)
    );

    pangolin::View &d_cam = pangolin::CreateDisplay()
            .SetBounds(0.0, 1.0, pangolin::Attach::Pix(175), 1.0, -1024.0f / 768.0f)
            .SetHandler(new pangolin::Handler3D(s_cam));


    while (pangolin::ShouldQuit() == false) {
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

        d_cam.Activate(s_cam);
        glClearColor(1.0f, 1.0f, 1.0f, 1.0f);

        glLineWidth(2);
        for (size_t i = 0; i < poses_est.size() - 1; i++) {
            glColor3f(1 - (float) i / poses_est.size(), 0.0f, (float) i / poses_est.size());
            glBegin(GL_LINES);
            auto p1 = poses_est[i], p2 = poses_est[i + 1];
            glVertex3d(p1.translation()[0], p1.translation()[1], p1.translation()[2]);
            glVertex3d(p2.translation()[0], p2.translation()[1], p2.translation()[2]);
            glEnd();
        }

	    for (size_t i = 0; i < poses_gt.size() - 1; i++) {
            // glColor3f(1 - (float) i / poses_gt.size(), 0.0f, (float) i / poses_gt.size());
            glColor3f(0.0, 0.0, 0.0);
            glBegin(GL_LINES);
            auto p1 = poses_gt[i], p2 = poses_gt[i + 1];
            glVertex3d(p1.translation()[0], p1.translation()[1], p1.translation()[2]);
            glVertex3d(p2.translation()[0], p2.translation()[1], p2.translation()[2]);
            glEnd();
        }
        pangolin::FinishFrame();
        usleep(5000);   // sleep 5 ms
    }

}

